{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "10ac4645-21b8-418a-bdfa-d8a82ec03afc",
   "metadata": {},
   "source": [
    "# 复现：RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval\n",
    "\n",
    "> RAG重要论文：递归抽象处理树形组织检索（https://metatwin.feishu.cn/wiki/KkM6wnDwgiDsJykVPPactqFknYb）\n",
    "- https://arxiv.org/abs/2401.18059\n",
    "- https://github.com/run-llama/llama_index/blob/main/llama-index-packs/llama-index-packs-raptor/examples/raptor.ipynb\n",
    "\n",
    "> 没有成功！！"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "93724ccc-1972-4a3b-ae99-e1d7cafe8b6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install llama-index llama-index-packs-raptor llama-index-vector-stores-qdrant -i https://pypi.tuna.tsinghua.edu.cn/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "1542b27b-f108-4e6c-942a-9d03ca08db34",
   "metadata": {},
   "outputs": [],
   "source": [
    "#!pip install llama-index-vector-stores-lancedb -i https://pypi.tuna.tsinghua.edu.cn/simple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "1589fff3-000b-44c2-88d7-dfc2d530fae6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.packs.raptor import RaptorPack\n",
    "\n",
    "#from llama_index.core.llama_pack import download_llama_pack\n",
    "#RaptorPack = download_llama_pack(\"RaptorPack\", \"./raptor_pack\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "b4429951-0699-46b9-bfa6-1db1c75806e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "id": "b3d6b7ef-ad2f-4650-91f8-251700c4d6ee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import SimpleDirectoryReader\n",
    "\n",
    "documents = SimpleDirectoryReader(input_files=[\"../datasets/2024-0103_O_editing.pdf\"]).load_data()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "3ae8a913-202b-4367-93f3-31bf92c1291b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Document(id_='d069606e-95d1-459b-b192-e1d14fad296b', embedding=None, metadata={'page_label': '1', 'file_name': '2024-0103_O_editing.pdf', 'file_path': '../datasets/2024-0103_O_editing.pdf', 'file_type': 'application/pdf', 'file_size': 608527, 'creation_date': '2024-05-23', 'last_modified_date': '2024-05-16'}, excluded_embed_metadata_keys=['file_name', 'file_type', 'file_size', 'creation_date', 'last_modified_date', 'last_accessed_date'], excluded_llm_metadata_keys=['file_name', 'file_type', 'file_size', 'creation_date', 'last_modified_date', 'last_accessed_date'], relationships={}, text='计算机集成制造系统  \\nComputer Integrated Manufacturing Systems  \\n \\n云模式下整车制造供应商选择与订单分配  \\n摘要：为了降低 云模式下整车制造零部件采购成本以及延迟交付风险，在考虑供应商零件价\\n格、零件质量、制造能力以及运输时间的基础上，引入供应商零件延迟送达惩罚函数，建立\\n了整车制造供应商选择与订单分配集成模型；设计遗传算法对模型进行求解，给出了遗传算\\n法的编码方式、交叉算子与变异算子。最后，以某型号轿车传动轴供应商选择与订单分配为\\n例，验证所建模型和求解算法的正确性和有效性。研究结果对于整车制 造各类零部件供应商\\n选择与订单分配具有重要的指导意义。  \\n关键词：云制造；供应商选择；订单分配；遗传算法  \\n中图分类号 ：TH166；TP391      文献标志码： A       \\nSupplier selection and order allocation for vehicle manufacturing \\nunder cloud mode  \\nAbstract:  In order to reduce the component procurement cost and delivery delay risk of vehicle \\nmanufacturing under cloud mode, a supplier selection and order allocation integrated model for \\nvehicle manufacturing was establishe d by introducing a supplier component delay delivery penalty \\nfunction, based on considering the supplier component price, quality, manufacturing capability and \\ntransportation time. A genetic algorithm was designed to solve the m odel, and the encoding method, \\ncrossover operator and muta tion operator of the genetic algorithm were given. Finally, taking the \\nsupplier selection and order allocation of a certain type of car transmission shaft as an example, the \\ncorrectness and effectiv eness of the proposed model and solution algorithm were verified. The \\nresearch results have important guiding significance for the supplier selection and order allocation \\nof various components of vehicle manufacturing.   \\nKeywords: cloud manufacturing; sup plier selection; order allocation; genetic algorithm  \\n ', start_char_idx=None, end_char_idx=None, text_template='{metadata_str}\\n\\n{content}', metadata_template='{key}: {value}', metadata_seperator='\\n')"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "documents[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "id": "c23d9196-5b8a-4e92-88e8-34a6e02e3a8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.node_parser import SentenceSplitter\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from llama_index.embeddings.openai import OpenAIEmbedding\n",
    "from llama_index.vector_stores.lancedb import LanceDBVectorStore\n",
    "from llama_index.core import VectorStoreIndex\n",
    "\n",
    "from llama_index.embeddings.ollama import OllamaEmbedding\n",
    "from llama_index.llms.ollama import Ollama\n",
    "\n",
    "from llama_index.core import Settings\n",
    "\n",
    "EMBEDDING_MODEL  = \"mixedbread-ai/mxbai-embed-large-v1\"\n",
    "GENERATION_MODEL = \"mistral-7b-v0.3\"\n",
    "\n",
    "# llm = MistralAI(model=GENERATION_MODEL)\n",
    "\n",
    "llm = Ollama(model=GENERATION_MODEL)\n",
    "\n",
    "ollama_embedding = OllamaEmbedding(\n",
    "    model_name= GENERATION_MODEL,\n",
    "    base_url=\"http://localhost:11434\",\n",
    "    ollama_additional_kwargs={\"mirostat\": 0},\n",
    ")\n",
    "\n",
    "\n",
    "Settings.llm = llm\n",
    "Settings.embed_model = ollama_embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "4f499368-4616-4936-8561-38d1890d6706",
   "metadata": {},
   "outputs": [],
   "source": [
    "vector_store = LanceDBVectorStore(uri=\"./tmp/lancedb\")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "index = VectorStoreIndex.from_documents(\n",
    "    documents, storage_context=storage_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "id": "7a670c78-f8a4-4248-80cf-2f9b814f2c96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating embeddings for level 0.\n",
      "Performing clustering for level 0.\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "OMP: Info #276: omp_set_nested routine deprecated, please use omp_set_max_active_levels instead.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Generating summaries for level 0 with 3 clusters.\n",
      "Level 0 created summaries/clusters: 3\n",
      "Generating embeddings for level 1.\n",
      "Performing clustering for level 1.\n",
      "Generating summaries for level 1 with 1 clusters.\n",
      "Level 1 created summaries/clusters: 1\n"
     ]
    },
    {
     "ename": "ArrowTypeError",
     "evalue": "struct fields don't match or are in the wrong order: Input fields: struct<_node_content: string, _node_type: string, doc_id: string, document_id: string, level: int64, parent_id: string, ref_doc_id: string> output fields: struct<_node_content: string, _node_type: string, creation_date: string, doc_id: string, document_id: string, file_name: string, file_path: string, file_size: int64, file_type: string, last_modified_date: string, page_label: string, ref_doc_id: string>",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mArrowTypeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[42], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m raptor_pack \u001b[38;5;241m=\u001b[39m \u001b[43mRaptorPack\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m      2\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdocuments\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m      3\u001b[0m \u001b[43m    \u001b[49m\u001b[43mvector_store\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvector_store\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# used for storage\u001b[39;49;00m\n\u001b[1;32m      4\u001b[0m \u001b[43m    \u001b[49m\u001b[43msimilarity_top_k\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# top k for each layer, or overall top-k for collapsed\u001b[39;49;00m\n\u001b[1;32m      5\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmode\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcollapsed\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# sets default mode\u001b[39;49;00m\n\u001b[1;32m      6\u001b[0m \u001b[43m    \u001b[49m\u001b[43mtransformations\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m[\u001b[49m\n\u001b[1;32m      7\u001b[0m \u001b[43m        \u001b[49m\u001b[43mSentenceSplitter\u001b[49m\u001b[43m(\u001b[49m\u001b[43mchunk_size\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m1000\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mchunk_overlap\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m50\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m      8\u001b[0m \u001b[43m    \u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m  \u001b[49m\u001b[38;5;66;43;03m# transformations applied for ingestion\u001b[39;49;00m\n\u001b[1;32m      9\u001b[0m \u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/packs/raptor/base.py:348\u001b[0m, in \u001b[0;36mRaptorPack.__init__\u001b[0;34m(self, documents, llm, embed_model, vector_store, similarity_top_k, mode, verbose, **kwargs)\u001b[0m\n\u001b[1;32m    336\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m__init__\u001b[39m(\n\u001b[1;32m    337\u001b[0m     \u001b[38;5;28mself\u001b[39m,\n\u001b[1;32m    338\u001b[0m     documents: List[BaseNode],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    345\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[1;32m    346\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    347\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Init params.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 348\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mretriever \u001b[38;5;241m=\u001b[39m \u001b[43mRaptorRetriever\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m    349\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdocuments\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    350\u001b[0m \u001b[43m        \u001b[49m\u001b[43membed_model\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43membed_model\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    351\u001b[0m \u001b[43m        \u001b[49m\u001b[43mllm\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mllm\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    352\u001b[0m \u001b[43m        \u001b[49m\u001b[43msimilarity_top_k\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msimilarity_top_k\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    353\u001b[0m \u001b[43m        \u001b[49m\u001b[43mvector_store\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mvector_store\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    354\u001b[0m \u001b[43m        \u001b[49m\u001b[43mmode\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mmode\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    355\u001b[0m \u001b[43m        \u001b[49m\u001b[43mverbose\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mverbose\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    356\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m    357\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/packs/raptor/base.py:139\u001b[0m, in \u001b[0;36mRaptorRetriever.__init__\u001b[0;34m(self, documents, tree_depth, similarity_top_k, llm, embed_model, vector_store, transformations, summary_module, existing_index, mode, **kwargs)\u001b[0m\n\u001b[1;32m    136\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39msimilarity_top_k \u001b[38;5;241m=\u001b[39m similarity_top_k\n\u001b[1;32m    138\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(documents) \u001b[38;5;241m>\u001b[39m \u001b[38;5;241m0\u001b[39m:\n\u001b[0;32m--> 139\u001b[0m     \u001b[43masyncio\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minsert\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdocuments\u001b[49m\u001b[43m)\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/nest_asyncio.py:30\u001b[0m, in \u001b[0;36m_patch_asyncio.<locals>.run\u001b[0;34m(main, debug)\u001b[0m\n\u001b[1;32m     28\u001b[0m task \u001b[38;5;241m=\u001b[39m asyncio\u001b[38;5;241m.\u001b[39mensure_future(main)\n\u001b[1;32m     29\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m---> 30\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mloop\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mrun_until_complete\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtask\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     31\u001b[0m \u001b[38;5;28;01mfinally\u001b[39;00m:\n\u001b[1;32m     32\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m task\u001b[38;5;241m.\u001b[39mdone():\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/nest_asyncio.py:98\u001b[0m, in \u001b[0;36m_patch_loop.<locals>.run_until_complete\u001b[0;34m(self, future)\u001b[0m\n\u001b[1;32m     95\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m f\u001b[38;5;241m.\u001b[39mdone():\n\u001b[1;32m     96\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mRuntimeError\u001b[39;00m(\n\u001b[1;32m     97\u001b[0m         \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mEvent loop stopped before Future completed.\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[0;32m---> 98\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mf\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mresult\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/asyncio/futures.py:201\u001b[0m, in \u001b[0;36mFuture.result\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    199\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__log_traceback \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[1;32m    200\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m--> 201\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception\u001b[38;5;241m.\u001b[39mwith_traceback(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exception_tb)\n\u001b[1;32m    202\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_result\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/asyncio/tasks.py:232\u001b[0m, in \u001b[0;36mTask.__step\u001b[0;34m(***failed resolving arguments***)\u001b[0m\n\u001b[1;32m    228\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m    229\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m exc \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    230\u001b[0m         \u001b[38;5;66;03m# We use the `send` method directly, because coroutines\u001b[39;00m\n\u001b[1;32m    231\u001b[0m         \u001b[38;5;66;03m# don't have `__iter__` and `__next__` methods.\u001b[39;00m\n\u001b[0;32m--> 232\u001b[0m         result \u001b[38;5;241m=\u001b[39m \u001b[43mcoro\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43msend\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43;01mNone\u001b[39;49;00m\u001b[43m)\u001b[49m\n\u001b[1;32m    233\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    234\u001b[0m         result \u001b[38;5;241m=\u001b[39m coro\u001b[38;5;241m.\u001b[39mthrow(exc)\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/packs/raptor/base.py:225\u001b[0m, in \u001b[0;36mRaptorRetriever.insert\u001b[0;34m(self, documents)\u001b[0m\n\u001b[1;32m    222\u001b[0m         node\u001b[38;5;241m.\u001b[39membedding \u001b[38;5;241m=\u001b[39m id_to_embedding[node\u001b[38;5;241m.\u001b[39mid_]\n\u001b[1;32m    223\u001b[0m         nodes_with_embeddings\u001b[38;5;241m.\u001b[39mappend(node)\n\u001b[0;32m--> 225\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mindex\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43minsert_nodes\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnodes_with_embeddings\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    227\u001b[0m \u001b[38;5;66;03m# set the current nodes to the new nodes\u001b[39;00m\n\u001b[1;32m    228\u001b[0m cur_nodes \u001b[38;5;241m=\u001b[39m new_nodes\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/core/indices/vector_store/base.py:330\u001b[0m, in \u001b[0;36mVectorStoreIndex.insert_nodes\u001b[0;34m(self, nodes, **insert_kwargs)\u001b[0m\n\u001b[1;32m    327\u001b[0m             node\u001b[38;5;241m.\u001b[39mobj \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m    329\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_callback_manager\u001b[38;5;241m.\u001b[39mas_trace(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minsert_nodes\u001b[39m\u001b[38;5;124m\"\u001b[39m):\n\u001b[0;32m--> 330\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_insert\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnodes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minsert_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    331\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_storage_context\u001b[38;5;241m.\u001b[39mindex_store\u001b[38;5;241m.\u001b[39madd_index_struct(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_index_struct)\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/core/indices/vector_store/base.py:312\u001b[0m, in \u001b[0;36mVectorStoreIndex._insert\u001b[0;34m(self, nodes, **insert_kwargs)\u001b[0m\n\u001b[1;32m    310\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_insert\u001b[39m(\u001b[38;5;28mself\u001b[39m, nodes: Sequence[BaseNode], \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39minsert_kwargs: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m    311\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Insert a document.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 312\u001b[0m     \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_add_nodes_to_index\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_index_struct\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mnodes\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minsert_kwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/core/indices/vector_store/base.py:234\u001b[0m, in \u001b[0;36mVectorStoreIndex._add_nodes_to_index\u001b[0;34m(self, index_struct, nodes, show_progress, **insert_kwargs)\u001b[0m\n\u001b[1;32m    232\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m nodes_batch \u001b[38;5;129;01min\u001b[39;00m iter_batch(nodes, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_insert_batch_size):\n\u001b[1;32m    233\u001b[0m     nodes_batch \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_get_node_with_embedding(nodes_batch, show_progress)\n\u001b[0;32m--> 234\u001b[0m     new_ids \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_vector_store\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd\u001b[49m\u001b[43m(\u001b[49m\u001b[43mnodes_batch\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43minsert_kwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    236\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_vector_store\u001b[38;5;241m.\u001b[39mstores_text \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_store_nodes_override:\n\u001b[1;32m    237\u001b[0m         \u001b[38;5;66;03m# NOTE: if the vector store doesn't store text,\u001b[39;00m\n\u001b[1;32m    238\u001b[0m         \u001b[38;5;66;03m# we need to add the nodes to the index struct and document store\u001b[39;00m\n\u001b[1;32m    239\u001b[0m         \u001b[38;5;28;01mfor\u001b[39;00m node, new_id \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mzip\u001b[39m(nodes_batch, new_ids):\n\u001b[1;32m    240\u001b[0m             \u001b[38;5;66;03m# NOTE: remove embedding from node to avoid duplication\u001b[39;00m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/llama_index/vector_stores/lancedb/base.py:178\u001b[0m, in \u001b[0;36mLanceDBVectorStore.add\u001b[0;34m(self, nodes, **add_kwargs)\u001b[0m\n\u001b[1;32m    176\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtable_name \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_connection\u001b[38;5;241m.\u001b[39mtable_names():\n\u001b[1;32m    177\u001b[0m     tbl \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_connection\u001b[38;5;241m.\u001b[39mopen_table(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtable_name)\n\u001b[0;32m--> 178\u001b[0m     \u001b[43mtbl\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43madd\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m    179\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    180\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_connection\u001b[38;5;241m.\u001b[39mcreate_table(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mtable_name, data)\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/lancedb/table.py:1107\u001b[0m, in \u001b[0;36mLanceTable.add\u001b[0;34m(self, data, mode, on_bad_vectors, fill_value)\u001b[0m\n\u001b[1;32m   1083\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Add data to the table.\u001b[39;00m\n\u001b[1;32m   1084\u001b[0m \u001b[38;5;124;03mIf vector columns are missing and the table\u001b[39;00m\n\u001b[1;32m   1085\u001b[0m \u001b[38;5;124;03mhas embedding functions, then the vector columns\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m   1104\u001b[0m \u001b[38;5;124;03m    The number of vectors in the table.\u001b[39;00m\n\u001b[1;32m   1105\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[1;32m   1106\u001b[0m \u001b[38;5;66;03m# TODO: manage table listing and metadata separately\u001b[39;00m\n\u001b[0;32m-> 1107\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[43m_sanitize_data\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1108\u001b[0m \u001b[43m    \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1109\u001b[0m \u001b[43m    \u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1110\u001b[0m \u001b[43m    \u001b[49m\u001b[43mmetadata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mmetadata\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1111\u001b[0m \u001b[43m    \u001b[49m\u001b[43mon_bad_vectors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mon_bad_vectors\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1112\u001b[0m \u001b[43m    \u001b[49m\u001b[43mfill_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfill_value\u001b[49m\u001b[43m,\u001b[49m\n\u001b[1;32m   1113\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1114\u001b[0m \u001b[38;5;66;03m# Access the dataset_mut property to ensure that the dataset is mutable.\u001b[39;00m\n\u001b[1;32m   1115\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_ref\u001b[38;5;241m.\u001b[39mdataset_mut\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/lancedb/table.py:89\u001b[0m, in \u001b[0;36m_sanitize_data\u001b[0;34m(data, schema, metadata, on_bad_vectors, fill_value)\u001b[0m\n\u001b[1;32m     87\u001b[0m         metadata\u001b[38;5;241m.\u001b[39mupdate(data\u001b[38;5;241m.\u001b[39mschema\u001b[38;5;241m.\u001b[39mmetadata \u001b[38;5;129;01mor\u001b[39;00m {})\n\u001b[1;32m     88\u001b[0m         data \u001b[38;5;241m=\u001b[39m data\u001b[38;5;241m.\u001b[39mreplace_schema_metadata(metadata)\n\u001b[0;32m---> 89\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[43m_sanitize_schema\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m     90\u001b[0m \u001b[43m        \u001b[49m\u001b[43mdata\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mschema\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mon_bad_vectors\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mon_bad_vectors\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mfill_value\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mfill_value\u001b[49m\n\u001b[1;32m     91\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m     92\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28misinstance\u001b[39m(data, Iterable):\n\u001b[1;32m     93\u001b[0m     data \u001b[38;5;241m=\u001b[39m _to_record_batch_generator(\n\u001b[1;32m     94\u001b[0m         data, schema, metadata, on_bad_vectors, fill_value\n\u001b[1;32m     95\u001b[0m     )\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/lancedb/table.py:1584\u001b[0m, in \u001b[0;36m_sanitize_schema\u001b[0;34m(data, schema, on_bad_vectors, fill_value)\u001b[0m\n\u001b[1;32m   1575\u001b[0m         \u001b[38;5;28;01mif\u001b[39;00m field\u001b[38;5;241m.\u001b[39mname \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mcolumn_names \u001b[38;5;129;01mand\u001b[39;00m (\n\u001b[1;32m   1576\u001b[0m             likely_vector_col \u001b[38;5;129;01mor\u001b[39;00m is_default_vector_col\n\u001b[1;32m   1577\u001b[0m         ):\n\u001b[1;32m   1578\u001b[0m             data \u001b[38;5;241m=\u001b[39m _sanitize_vector_column(\n\u001b[1;32m   1579\u001b[0m                 data,\n\u001b[1;32m   1580\u001b[0m                 vector_column_name\u001b[38;5;241m=\u001b[39mfield\u001b[38;5;241m.\u001b[39mname,\n\u001b[1;32m   1581\u001b[0m                 on_bad_vectors\u001b[38;5;241m=\u001b[39mon_bad_vectors,\n\u001b[1;32m   1582\u001b[0m                 fill_value\u001b[38;5;241m=\u001b[39mfill_value,\n\u001b[1;32m   1583\u001b[0m             )\n\u001b[0;32m-> 1584\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mpa\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mTable\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_arrays\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m   1585\u001b[0m \u001b[43m        \u001b[49m\u001b[43m[\u001b[49m\u001b[43mdata\u001b[49m\u001b[43m[\u001b[49m\u001b[43mname\u001b[49m\u001b[43m]\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;28;43;01mfor\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mname\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;129;43;01min\u001b[39;49;00m\u001b[43m \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mnames\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mschema\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mschema\u001b[49m\n\u001b[1;32m   1586\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m   1588\u001b[0m \u001b[38;5;66;03m# just check the vector column\u001b[39;00m\n\u001b[1;32m   1589\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m VECTOR_COLUMN_NAME \u001b[38;5;129;01min\u001b[39;00m data\u001b[38;5;241m.\u001b[39mcolumn_names:\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/table.pxi:4625\u001b[0m, in \u001b[0;36mpyarrow.lib.Table.from_arrays\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/table.pxi:1562\u001b[0m, in \u001b[0;36mpyarrow.lib._sanitize_arrays\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/array.pxi:382\u001b[0m, in \u001b[0;36mpyarrow.lib.asarray\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/table.pxi:574\u001b[0m, in \u001b[0;36mpyarrow.lib.ChunkedArray.cast\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/compute.py:404\u001b[0m, in \u001b[0;36mcast\u001b[0;34m(arr, target_type, safe, options, memory_pool)\u001b[0m\n\u001b[1;32m    402\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m    403\u001b[0m         options \u001b[38;5;241m=\u001b[39m CastOptions\u001b[38;5;241m.\u001b[39msafe(target_type)\n\u001b[0;32m--> 404\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mcall_function\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[38;5;124;43mcast\u001b[39;49m\u001b[38;5;124;43m\"\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m[\u001b[49m\u001b[43marr\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43moptions\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmemory_pool\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/_compute.pyx:590\u001b[0m, in \u001b[0;36mpyarrow._compute.call_function\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/_compute.pyx:385\u001b[0m, in \u001b[0;36mpyarrow._compute.Function.call\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/error.pxi:154\u001b[0m, in \u001b[0;36mpyarrow.lib.pyarrow_internal_check_status\u001b[0;34m()\u001b[0m\n",
      "File \u001b[0;32m~/opt/anaconda3/envs/llama-index/lib/python3.10/site-packages/pyarrow/error.pxi:91\u001b[0m, in \u001b[0;36mpyarrow.lib.check_status\u001b[0;34m()\u001b[0m\n",
      "\u001b[0;31mArrowTypeError\u001b[0m: struct fields don't match or are in the wrong order: Input fields: struct<_node_content: string, _node_type: string, doc_id: string, document_id: string, level: int64, parent_id: string, ref_doc_id: string> output fields: struct<_node_content: string, _node_type: string, creation_date: string, doc_id: string, document_id: string, file_name: string, file_path: string, file_size: int64, file_type: string, last_modified_date: string, page_label: string, ref_doc_id: string>"
     ]
    }
   ],
   "source": [
    "raptor_pack = RaptorPack(\n",
    "    documents,\n",
    "    vector_store=vector_store,  # used for storage\n",
    "    similarity_top_k=2,  # top k for each layer, or overall top-k for collapsed\n",
    "    mode=\"collapsed\",  # sets default mode\n",
    "    transformations=[\n",
    "        SentenceSplitter(chunk_size=1000, chunk_overlap=50)\n",
    "    ],  # transformations applied for ingestion\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "09196079-b573-4373-8de4-d1e7f6ab53b4",
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'raptor_pack' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[37], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m nodes \u001b[38;5;241m=\u001b[39m \u001b[43mraptor_pack\u001b[49m\u001b[38;5;241m.\u001b[39mrun(\n\u001b[1;32m      2\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m软件安全流程有哪些？\u001b[39m\u001b[38;5;124m\"\u001b[39m, mode\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mtree_traversal\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[1;32m      3\u001b[0m )\n",
      "\u001b[0;31mNameError\u001b[0m: name 'raptor_pack' is not defined"
     ]
    }
   ],
   "source": [
    "nodes = raptor_pack.run(\n",
    "    \"软件安全流程有哪些？\", mode=\"tree_traversal\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2533acc-fe3b-43ba-b1b2-487ee24b4253",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
